package com.kunlun.common.utils;

import com.kunlun.common.model.TreeModel;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang.StringUtils;
import org.springframework.util.ObjectUtils;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class TreeUtil {

    public static <T extends TreeModel> List<T> buildTree(List<T> list) {
        List<T> rootNoes = new ArrayList<>();
        Map<String, List<T>> parentMap = new HashMap<>();
        prepareTreeNode(list, rootNoes, parentMap);
        buildTreeNode(rootNoes, parentMap);
        return rootNoes;
    }

    private static <T extends TreeModel> void buildTreeNode(List<T> rootNoes, Map<String, List<T>> parentMap) {
        for (T rootNoe : rootNoes) {
            List<T> childrens = parentMap.get(rootNoe.getId());
            if (CollectionUtils.isNotEmpty(childrens)) {
                rootNoe.setChildren(childrens);
                buildTreeNode(childrens, parentMap);
            }
        }
    }

    private static <T extends TreeModel> void prepareTreeNode(List<T> list, List<T> rootNoes, Map<String, List<T>> parentMap) {
        for (T t : list) {
            if (StringUtils.isEmpty(t.getParentId())) {
                rootNoes.add(t);
            } else {
                List<T> values = parentMap.get(t.getParentId());
                if (ObjectUtils.isEmpty(values)) {
                    values = new ArrayList<>();
                    values.add(t);
                    parentMap.put(t.getParentId(), values);
                } else {
                    String longCode1 = distillParentLongCode(values.get(0).getLongCode());
                    String longCode2 = distillParentLongCode(t.getLongCode());
                    if (longCode1.equals(longCode2)) {
                        values.add(t);
                    }
                }
            }
        }
    }

    private static String distillParentLongCode(String longCode) {
        int index = longCode.lastIndexOf("_");
        return longCode.substring(0, index);
    }

    public static <T extends TreeModel> void addTreeNode(List<T> companyModels, Map<String, List<T>> departmentMap) {
        for (T t : companyModels) {
            List<T> children = t.getChildren();
            if (CollectionUtils.isNotEmpty(children)) {
                addTreeNode(children, departmentMap);
            }

            List<T> departmentModels = departmentMap.get(t.getLongCode());
            if (CollectionUtils.isNotEmpty(departmentModels)) {
                for (T departmentModel : departmentModels) {
                    if (StringUtils.isEmpty(t.getParentId())) {
                        departmentModel.setLongCode(t.getId());
                    } else {
                        departmentModel.setLongCode(t.getLongCode() + "_" + departmentModel.getId());
                    }
                }
                if (CollectionUtils.isNotEmpty(children)) {
                    t.getChildren().addAll(departmentModels);
                } else {
                    t.setChildren(departmentModels);
                }
            }
        }
    }
}
